package com.lemon.cloud.gateway.filter;

import cn.hutool.core.util.StrUtil;
import com.lemon.cloud.gateway.config.GatewayConfigProperties;
import lombok.extern.slf4j.Slf4j;
import org.jsoup.Jsoup;
import org.jsoup.safety.Safelist;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.AntPathMatcher;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * xss防御
 */
@Component
@Slf4j
public class XssRequestFilter implements GlobalFilter, Ordered {

    @Autowired
    private GatewayConfigProperties configProperties;

    private final AntPathMatcher PATH_MATCHER = new AntPathMatcher();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        if (!enableXss(request)) {
            return chain.filter(exchange);
        }
        return DataBufferUtils.join(request.getBody())
                .defaultIfEmpty(toDataBuffer(StrUtil.EMPTY.getBytes()))
                .flatMap(dataBuffer -> {
                    // 取出body中的参数
                    byte[] oldBytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(oldBytes);
                    String bodyString = new String(oldBytes, StandardCharsets.UTF_8);
                    bodyString = Jsoup.clean(bodyString, Safelist.basicWithImages());
                    ServerHttpRequest newRequest = request.mutate().uri(request.getURI()).build();
                    // 重新构造body
                    byte[] newBytes = bodyString.getBytes(StandardCharsets.UTF_8);
                    DataBuffer bodyDataBuffer = toDataBuffer(newBytes);
                    Flux<DataBuffer> bodyFlux = Flux.just(bodyDataBuffer);
                    // 重新构造header
                    HttpHeaders headers = new HttpHeaders();
                    headers.putAll(request.getHeaders());
                    int length = newBytes.length;
                    headers.remove(HttpHeaders.CONTENT_LENGTH);
                    headers.setContentLength(length);
                    newRequest = new ServerHttpRequestDecorator(newRequest) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return bodyFlux;
                        }

                        @Override
                        public HttpHeaders getHeaders() {
                            return headers;
                        }
                    };
                    return chain.filter(exchange.mutate().request(newRequest).build());
                });
    }

    @Override
    public int getOrder() {
        return 1;
    }

    private DataBuffer toDataBuffer(byte[] bytes) {
        DefaultDataBufferFactory defaultDataBufferFactory = new DefaultDataBufferFactory();
        DataBuffer buffer = defaultDataBufferFactory.allocateBuffer(bytes.length);
        buffer.write(bytes);
        return buffer;
    }

    /**
     * 校验请求参数
     *
     * @param request
     */
    public boolean enableXss(ServerHttpRequest request) {
        if (request.getMethod() != null && StrUtil.equalsAnyIgnoreCase(request.getMethod().name(), HttpMethod.POST.name(), HttpMethod.PUT.name())) {
            //校验xss
            URI uri = request.getURI();
            List<String> excludeXssPatterns = configProperties.getExcludeXssPatterns();
            request.getHeaders();
            for (String e : excludeXssPatterns) {
                if (!PATH_MATCHER.match(e, uri.getPath())) {
                    return true;
                }
            }
        }
        return false;
    }

}
